from ultralytics import YOLO
import cv2
import argparse

# 创建 ArgumentParser 对象
parser = argparse.ArgumentParser(description='输入输出文件')

# 添加参数
parser.add_argument('-f', help='输入文件')
parser.add_argument('-d', help='输出文件', default='default value')


# Load a model
#model = YOLO("yolo11n.pt")

def model_train():
    # Train the model
    print("start to train")
    train_results = model.train(
        data="coco8.yaml",  # path to dataset YAML
        epochs=100,  # number of training epochs
        imgsz=640,  # training image size
        device="cpu",  # device to run on, i.e. device=0 or device=0,1,2,3 or device=cpu
    )
    print(train_results)

def model_val():
    # Evaluate model performance on the validation set
    metrics = model.val()


def model_inference(image):
    # Perform object detection on an image
    results = model(image)
    print(results)

def model_export():
    # Export the model to ONNX format
    return  model.export(format="onnx")  # return path to exported model

def predict(chosen_model, img, classes=[], conf=0.5):
    if classes:
        results = chosen_model.predict(img, classes=classes, conf=conf)
    else:
        results = chosen_model.predict(img, conf=conf)
    return results

def predict_and_detect(chosen_model, img, classes=[], conf=0.5, rectangle_thickness=2, text_thickness=1):
    results = predict(chosen_model, img, classes, conf=conf)
    for result in results:
        for box in result.boxes:
            cv2.rectangle(img, (int(box.xyxy[0][0]), int(box.xyxy[0][1])),
                          (int(box.xyxy[0][2]), int(box.xyxy[0][3])), (255, 0, 0), rectangle_thickness)
            cv2.putText(img, f"{result.names[int(box.cls[0])]}",
                        (int(box.xyxy[0][0]), int(box.xyxy[0][1]) - 10),
                        cv2.FONT_HERSHEY_PLAIN, 1, (255, 0, 0), text_thickness)
    return img, results


def gen_detect(filein, fileout):
    print("file %s to %s" % (filein, fileout))
    image = cv2.imread(filein)
    result_img, _ = predict_and_detect(model, image, classes=[], conf=0.5)
    cv2.imwrite(fileout, result_img)

if __name__ == '__main__':
    args = parser.parse_args()
    
    model = YOLO("outmodels/best.onnx")
    gen_detect(args.f, args.d)

    # model_train()

# python yolov11-detect.py -f /Users/v/Documents/proj/testdata/origin_imgs/images_labelimg_yolo/IMG_20181201_181804.jpg -d /Users/v/Documents/proj/testdata/origin_imgs/r1-out.jpg

